import os
import sys
script_dir = os.path.dirname(os.path.realpath(__file__))
parent_dir = os.path.dirname(os.path.dirname(script_dir))
sys.path.append(parent_dir)

import torch
import torch.nn as nn
import torch.nn.functional as F

from DeepPPG.model.resnet_t_dist import MetaResNet1d

def save_best_model(model_path, save_path, model_name):
    model = torch.load(model_path)
    model.eval()

    print("model loaded.")
    torch.save(model.state_dict(), os.path.join(save_path, f"{model_name}_state_dict.pth"))
    print("model saved.")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', type=str, default="seq2one_MetaResNet_log_0.01-500-criterion_FocalLoss-criterion_hp__'alpha'_ [0.25, 0.75], 'gamma'_ 2_-8-True-0.2-12-0.01-AF_1year-0.0005-0_models/best_model.pt", help='Path to the model file')
    parser.add_argument('--save_path', type=str, default="DeepECG/performance_test/save_best_models/1_year_meta_td_patient_level_all_lead/", help='Path to save the best model')
    parser.add_argument('--model_name', type=str, default="MetaResNet", help='Name of the model')
    args = parser.parse_args()

    if not os.path.exists(args.save_path):
        os.makedirs(args.save_path)

    save_best_model(args.model_path, args.save_path, args.model_name)

    # Test loading model
    if args.model_name == "MetaResNet":
        model = MetaResNet1d(nOUT=2, in_ch=12, out_ch=256, mid_ch=64, if_t_dist=True)
        # model = MetaResNet1d(nOUT=1, in_ch=1, out_ch=256, mid_ch=64)
    
    model.load_state_dict(torch.load(os.path.join(args.save_path, f"{args.model_name}_state_dict.pth")))
    model.eval()
    print("Model loaded successfully!")

